/*
 * decode image using different context, by edhong
 * Jan-March 2008
 *
 * 
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
 *
 */


#include "SubbandBits.h"
#include "Characteristics.h"
#include "coder.h"
#include <math.h>


#define USG_STRING "[-bp %d:numbitplanes] [-bt %d:numbits]  [-cf %s:ctdfilename] [-bs %d:bitsperiteration; 1 = by bitplane] [-i %s:initialimagename] [-d %d:debugflag=0] %s:bitsfile %s:output_imgname"

QccString initialimagename = "";

QccWAVWavelet Wavelet;
QccString WaveletFilename = QCCWAVWAVELET_DEFAULT_WAVELET;
QccString Boundary = "symmetric";
QccString ctdfilename = "";

int NumSubbands;
int NumBitplanes = -1;
int numbits = -1; 
int bitsperiteration = -1;
int CharacteristicType;
int ColorCharacteristicType;
int debugflag = 0;
int colortype;

int ImageNumCols;
int ImageNumRows;
int sortNum;

QccIMGImage OriginalImage;
QccIMGImage OutputImage;
QccIMGImageComponent temp1C;
QccIMGImageComponent temp2C;
QccWAVSubbandPyramid SubbandPyramid[3];
QccWAVSubbandPyramid OriginalSubbandPyramid[3];
QccBitBuffer input_buffer;
QccENTArithmeticModel *model = NULL; 


void RoundAndTruncate(QccWAVSubbandPyramid *sbp) {
  int i,j;
  for (i=0; i < sbp->num_rows; i++) {
    for (j=0; j < sbp->num_cols; j++) 
       sbp->matrix[i][j] =QccIMGImageComponentClipPixel(round(sbp->matrix[i][j]));
  }
}
 

void ConvertToRGBRoundAndTruncateConvertBack(QccWAVSubbandPyramid *sbp, int ctype) {
  int i,j,k, R, G, B;
  switch (ctype) {
    case QCCIMGCOLORTYPE_YCbCr:
      for (i=0; i < sbp->num_rows; i++) {
        for (j=0; j < sbp->num_cols; j++) {
                QccIMGImageJPEG_YCbCrtoRGB(&R, &G, &B, 
                  sbp[0].matrix[i][j], 
                  sbp[1].matrix[i][j], 
                  sbp[2].matrix[i][j] );
                QccIMGImageRGBtoJPEG_YCbCr(
                   &sbp[0].matrix[i][j], 
                   &sbp[1].matrix[i][j], 
                   &sbp[2].matrix[i][j], 
                   R, G, B);
        }
       }
       break;                        
     case QCCIMGCOLORTYPE_YUV:
       for (i=0; i < sbp->num_rows; i++) {
         for (j=0; j < sbp->num_cols; j++) {
                QccIMGImageYUVtoRGB(&R, &G, &B, 
                  sbp[0].matrix[i][j], 
                  sbp[1].matrix[i][j], 
                  sbp[2].matrix[i][j] );
                QccIMGImageRGBtoYUV(
                   &sbp[0].matrix[i][j], 
                   &sbp[1].matrix[i][j], 
                   &sbp[2].matrix[i][j], 
                   R, G, B);
          }
        }
        break;
      case QCCIMGCOLORTYPE_RGB:
         for (k=0; k < 3; k++) {
              RoundAndTruncate(&sbp[k]);
         }
         break;
      default: 
         printf("Unimplemented color type %d\n", colortype);                   
    }
}


int main(int argc, char *argv[])
{
  unsigned long long flags[2];
  
  QccInit(argc, argv);

  int i,j,k;
  QccWAVWaveletInitialize(&Wavelet);
  QccIMGImageInitialize(&OriginalImage);
  QccIMGImageInitialize(&OutputImage);
  QccIMGImageComponentInitialize(&temp1C);
  QccIMGImageComponentInitialize(&temp2C);
  for (i=0; i < 3; i++) {
    QccWAVSubbandPyramidInitialize(&SubbandPyramid[i]);
    QccWAVSubbandPyramidInitialize(&OriginalSubbandPyramid[i]);
  }
  QccBitBufferInitialize(&input_buffer);
  
  if (QccParseParameters(argc, argv,
                         USG_STRING,
                         &NumBitplanes,
			 &numbits,
			 &ctdfilename,
			 &bitsperiteration,
			 OriginalImage.filename,
			 &debugflag,
                         input_buffer.filename,
                         OutputImage.filename))
    QccErrorExit();

  if (NumBitplanes >0  && numbits >0)
    {
      QccErrorAddMessage("%s: Can only specify one positve number in either bitplanes and numbits",
                         argv[0]);
      QccErrorExit();

    }

    
  if (QccWAVWaveletCreate(&Wavelet, WaveletFilename, Boundary))
    {
      QccErrorAddMessage("%s: Error calling QccWAVWaveletCreate()",
                         argv[0]);
      QccErrorExit();
    }
 
  if (!QccFileExists(input_buffer.filename)) {
    QccErrorAddMessage("%s: cannot find file %s ",
		       argv[0], input_buffer.filename);
    QccErrorExit();
  }

  input_buffer.type = QCCBITBUFFER_INPUT; 
  if(QccBitBufferStart(&input_buffer)) {
    QccErrorAddMessage("%s:  Error calling QccBitBufferStart",
		       argv[0]);
    QccErrorExit();
  }


// determine color/not color by output filename
  if (QccIMGImageDetermineType(&OutputImage)) {
    QccErrorAddMessage("%s: Error calling QccIMGImageDetermineType() on  %s",
		       argv[0], OutputImage.filename);
    QccErrorExit();
  }

  // initial stuff to get from bitbuffer: numcols, numRow, 
  //  log_2 of initial threshold (as int0, CharacteristicType,
  //  sorting method (4 bits), then actualnumbitplanes (5 bits)
  //   colortype (if necessary, 2 bits), 
  //   colorCharacteristicType (if necessary, 4 bits), 
//   encodectd, and  <encoding of ctd> if necessary. 

  int numbitsimagelength = 14;
  int error = 0;
  int ctd_encoded;
  int actualnumbitplanes;

  int initialthreshpower; 
  error= QccBitBufferGetBits(&input_buffer, &ImageNumCols, numbitsimagelength);
  error+=QccBitBufferGetBits(&input_buffer, &ImageNumRows, numbitsimagelength);
  error+=QccBitBufferGetBits(&input_buffer, &initialthreshpower, 4);
  error+=QccBitBufferGetBits(&input_buffer, &CharacteristicType, 8);
  error+=QccBitBufferGetBits(&input_buffer, &sortNum, 4);
  error+=QccBitBufferGetBits(&input_buffer, &actualnumbitplanes, 5);
  if (QccIMGImageColor(&OutputImage)) {
    error+=QccBitBufferGetBits(&input_buffer, &colortype, 2);
    error += QccBitBufferGetBits(&input_buffer, &ColorCharacteristicType, 4);
  }
  error+=QccBitBufferGetBit(&input_buffer, &ctd_encoded);
  if (error >0) {
    printf( "(%s): Error calling QccBitBufferGetBit(s)\n", argv[0]); 
    return 1;
  }

  int NumLevels = computeNumWaveletLevels(ImageNumRows, ImageNumCols);
  NumSubbands = QccWAVSubbandPyramidNumLevelsToNumSubbands(NumLevels);
  InitAllC(NumSubbands, NumBitplanes);

  if (getCharacteristicSet(CharacteristicType, flags ))
    {
      QccErrorExit();
    }
  if (QccIMGImageColor(&OutputImage)) {
     if (addColorCharacteristicSet(ColorCharacteristicType, flags))
     {
       QccErrorExit();
     }
  }
  contexttree ct;
  InitAllC(NumSubbands, actualnumbitplanes);

  if (ctd_encoded) {
    if (ctdfilename[0]!='\0') {
      QccErrorAddMessage("%s:  Cannot include ctd filename, ctd already encoded in bitstream",
			 argv[0]);
      QccErrorExit();
    }
    // decode ctd tree from input_buffer here. 
    if (decodectd_bitbuffer(&input_buffer, flags, &ct)) {
       QccErrorAddMessage("%s: Error calling decodectd_bitbuffer()",
			 argv[0]);
        QccErrorExit();

     }
  }
  else {
    if (ctdfilename[0]=='\0') {
      QccErrorAddMessage("%s:  must include ctd filename, ctd not encoded in bitstream",
			 argv[0]);
      QccErrorExit();
    }
    // read ctd  from input ctd file. 

     if (readctdfile(ctdfilename, flags, &ct)) {
       QccErrorAddMessage("%s: Error calling readctdfile()",
			 argv[0]);
        QccErrorExit();

     }
  }

    if (debugflag) {
      printf("Header (possibly incl. ctd tree)  used %d bits.\n", input_buffer.bit_cnt); 
    }


  if (debugflag==1) printcontexttree(&ct);

  int numComponents;
  numComponents = 1;


  if (QccIMGImageColor(&OutputImage))  {
    numComponents = 3;
    if  (QccIMGImageSetColorType(&OutputImage, colortype)) 
    {
      QccErrorAddMessage("%s: Error calling QccIMGImageSetColorType()",
                         argv[0]);
      QccErrorExit();
    
    }  
  }

  if (OriginalImage.filename[0] != '\0') {
    if (QccIMGImageColor(&OutputImage))  { // read in original image using correct color space.
      if  (QccIMGImageSetColorType(&OriginalImage, colortype)) 
      {
        QccErrorAddMessage("%s: Error calling QccIMGImageSetColorType()",
                         argv[0]);
        QccErrorExit();
    
      }
    }  
    if (QccIMGImageRead(&OriginalImage))
      {
	QccErrorAddMessage("%s: Error calling QccIMGImageRead() on original image %s",
			   argv[0], OriginalImage.filename);
	QccErrorExit();
      }

    if (OriginalImage.image_type != OutputImage.image_type) {

      QccErrorAddMessage("%s: Image type of original image and output image must match",
			 argv[0]);
      QccErrorExit();
    }

    QccIMGImageComponentCopy(&temp1C, &OriginalImage.Y);
    

    OriginalSubbandPyramid[0].num_levels = 0;
    OriginalSubbandPyramid[0].num_rows = temp1C.num_rows;
    OriginalSubbandPyramid[0].num_cols = temp1C.num_cols;
    OriginalSubbandPyramid[0].matrix = temp1C.image;
    
    if (OriginalSubbandPyramid[0].num_rows != ImageNumRows  
	||
	OriginalSubbandPyramid[0].num_cols !=  ImageNumCols) {
      QccErrorAddMessage("%s:  number of rows and cols  in original image do not match input file",
			 argv[0]);
      QccErrorExit();
    }

    if (numComponents >1) {
      OriginalSubbandPyramid[1] =       OriginalSubbandPyramid[0] ;
      OriginalSubbandPyramid[2] =       OriginalSubbandPyramid[0] ;
      temp1C.image = NULL;
      QccIMGImageComponentCopy(&temp1C, &OriginalImage.U);
      OriginalSubbandPyramid[1].matrix = temp1C.image;
      temp1C.image = NULL;
      QccIMGImageComponentCopy(&temp1C, &OriginalImage.V);
      OriginalSubbandPyramid[2].matrix = temp1C.image;
    }

    temp1C.num_rows = ImageNumRows;
    temp1C.num_cols = ImageNumCols;
    temp2C.num_rows = ImageNumRows;
    temp2C.num_cols = ImageNumCols;


  // subtract mean value of possible 8-bit values, for luminance plane. 
  int i,j,k;
  for (i=0; i < OriginalSubbandPyramid[0].num_rows; i++) {
    for (j=0; j < OriginalSubbandPyramid[0].num_cols; j++) 
       OriginalSubbandPyramid[0].matrix[i][j] -= (255.0/2);
  }

    for (k=1; k < numComponents; k++) 
      for (i=0; i < OriginalSubbandPyramid[k].num_rows; i++) {
        for (j=0; j < OriginalSubbandPyramid[k].num_cols; j++) 
            switch (colortype) {
                case  QCCIMGCOLORTYPE_YCbCr:
                 case QCCIMGCOLORTYPE_YUV:
                    // do nothing, ranges already correct (from -127.5 to 127.5)
                    break;
                case QCCIMGCOLORTYPE_RGB:
                    OriginalSubbandPyramid[k].matrix[i][j] -= (255.0/2);
                   break;
                default: 
                   printf("Unimplemented color type %d\n", colortype);                   
            }
      } 
  

    for (i=0; i < numComponents; i++) {
      if (QccWAVSubbandPyramidDWT(&OriginalSubbandPyramid[i],
				  NumLevels,
				  &Wavelet))
	{
	  QccErrorAddMessage("%s: Error calling QccWAVSubbandPyramidDWT()",
			     argv[0]);
	  QccErrorExit();
	}
    }
  }

  printf("Reading bitsfile %s, CharacteristicType %d, ",
	 input_buffer.filename, CharacteristicType);
  if (!ctd_encoded) 
    printf(" using ctdfile %s,", ctdfilename);
  printf(" writing to %s \n", OutputImage.filename);
  printf("  parameters: %d x %d image, ", ImageNumCols, ImageNumRows);
  printf ("%d bit-planes, ", NumBitplanes);
  printf("%d numbits, ", numbits);
  printf("%d threshold, ", 1<<initialthreshpower);
  printf("%d wavelet levels, sortNum = %d,  debug=%d\n", NumLevels,  sortNum, debugflag);
  printf("  %d bits per iteration,  OriginalImage:%s\n", bitsperiteration, OriginalImage.filename);
  if (numComponents==1) printf("Assuming grey-scale image\n");
  else printf("Decoding color image with colortype %d\n", colortype);



  int encode_eof = 0;  
  int numContexts;
  int *num_symbols;

  numContexts = ct.numleafnodes;
  num_symbols = (int *)malloc(sizeof(int)*numContexts);

  for ( i=0; i < numContexts; i++) {
    num_symbols[i] = 2;
  }

  model = QccENTArithmeticDecodeStart2(&input_buffer, 
				       num_symbols, 
				       numContexts, 
				       NULL, 
				       QCCENT_ANYNUMBITS,
				       encode_eof);

  if(model == NULL){
    QccErrorAddMessage("%s:  Error calling QccENTArithmeticDecodeStart2",
		       argv[0]);
    QccErrorExit();
  } 

 
  QccENTArithmeticSetModelAdaption(model, QCCENT_TKADAPTIVE);

  int totalbits;
  SubbandComponent sc[3];
  for (i=0; i < numComponents; i++) {
    SubbandPyramid[i].num_rows = ImageNumRows;
    SubbandPyramid[i].num_cols = ImageNumCols;
    SubbandPyramid[i].num_levels = NumLevels;
    QccWAVSubbandPyramidAlloc(&SubbandPyramid[i]);
  }

  // decode file. 
  int approxbitsused = 0;
  int result; 
  for ( i=0; i < numComponents; i++) {
    SubbandComponentInitialize(&sc[i], &SubbandPyramid[i]);
  }
  decode_subband_component_start(sc, numComponents, model, &input_buffer, &ct, 
				 get_sorting_function(sortNum),
				 NumBitplanes, numbits, &totalbits, debugflag);

  printf("Outputting:  approx_bits_used");
  if (OriginalImage.filename[0] != '\0') {
    printf(" SBP_MSE Image_MSE Image_MSE_after_rounding");
  }
  printf("\n");

  do {

    if (bitsperiteration==1) { 
      // read one bitplane by setting limit appropriately
      sbploc curr, curr2;
      curr = get_currentloc();
      //      printsbploc(&curr);
      decode_subband_component_update_limit(curr.k+1, numbits);
      result = decode_subband_component_continue(INT_MAX, debugflag);
      //      printf("result %d\n",result);
	curr2 = get_currentloc();
	//      printsbploc(&curr2);
      if ((curr2.k == curr.k)||
	  curr2.k == NumBitplanes)
	result = 1; // done. 
      else 
	result = 0;
    }
    else {
      int newapproxbitsused; 
      do {
       // repeat until at least 10 decoder bits have passed since last iteration. 
         result = decode_subband_component_continue(bitsperiteration, debugflag);
         newapproxbitsused = input_buffer.bit_cnt + model->garbage_bits; // total bits used by decoder. 
         newapproxbitsused -= (QCCENT_CODEVALUEBITS -1); // decoder always reads about this many bits ahead. 
      } while ((result==0) && (newapproxbitsused < approxbitsused+10)); 
      
    }
    //    printf("decode_subband_component_continue returned: %d\n", result);
    approxbitsused = input_buffer.bit_cnt + model->garbage_bits; // total bits used by decoder. 
    approxbitsused -= (QCCENT_CODEVALUEBITS -1); // decoder always reads about this many bits ahead. 
    
    printf(" %d", approxbitsused);
    
    if (debugflag ==2) {
      if (numComponents ==1) {
	printSubbandComponent(sc);
      }
      else {
	for (i=0; i < numComponents; i++) {
	  printf("**Component number %d\n", i);
	  printSubbandComponent(&sc[i]);
	}
      }
    }
    sbploc myloc;
    myloc = get_currentloc();

    if (debugflag) {
      printf("After %d bits: ", totalbits); printsbploc(&myloc);
    }
    // put in sc bits into subbandpyramid. 
    for (i=0; i < numComponents; i++) {
      SubbandPyramid[i].num_levels = NumLevels;
      copySubbandComponentToSubbandPyramid(&SubbandPyramid[i],&sc[i], 
					 &myloc, initialthreshpower);
    }
    // output image quality if original image was given. 
    if (OriginalImage.filename[0] != '\0') {
      double total; 
      total =0; 
      for (i=0; i < numComponents; i++) {
	temp1C.image = OriginalSubbandPyramid[i].matrix;
	temp2C.image = SubbandPyramid[i].matrix;
	total += 
	  QccIMGImageComponentMSE(&temp1C, &temp2C)/ numComponents;
      }
      printf(" %f", total);
    }
    // now invert  image file. 
    for ( i=0; i < numComponents; i++) 
      if (QccWAVSubbandPyramidInverseDWT(&SubbandPyramid[i],
					 &Wavelet))
	{
	  QccErrorAddMessage("%s: Error calling QccWAVSubbandPyramidInverseDWT()",
			     argv[0]);
	  QccErrorExit();
	}
    // add estimated mean back (luminance only)
  for (i=0; i < SubbandPyramid[0].num_rows; i++) {
    for (j=0; j < SubbandPyramid[0].num_cols; j++) 
       SubbandPyramid[0].matrix[i][j] += (255.0/2);
  }
    // go back to original range in chrominance.
    for (k=1; k < numComponents; k++) 
      for (i=0; i < SubbandPyramid[k].num_rows; i++) {
        for (j=0; j < SubbandPyramid[k].num_cols; j++) 
            switch (colortype) {
                case  QCCIMGCOLORTYPE_YCbCr:
                case  QCCIMGCOLORTYPE_YUV:
                    // do nothing. 
                   break;
                case QCCIMGCOLORTYPE_RGB:
                      SubbandPyramid[k].matrix[i][j] += (255.0/2);
                   break;
                default: 
                   printf("Unimplemented color type %d\n", colortype);                   
            }
      } 
  
    
    if (OriginalImage.filename[0] != '\0') { // compute ImageMSE (w/o rounding)
      double total; 
      temp1C.image = OriginalImage.Y.image;
      temp2C.image = SubbandPyramid[0].matrix;
      total = 	 QccIMGImageComponentMSE(&temp1C, &temp2C);
    //        printf(" %f", total);
      if (numComponents > 1) {
	temp1C.image = OriginalImage.U.image;
	temp2C.image = SubbandPyramid[1].matrix;
	total += QccIMGImageComponentMSE(&temp1C, &temp2C);
	//	printf(" %f", total);
	temp1C.image = OriginalImage.V.image;
	temp2C.image = SubbandPyramid[2].matrix;
	total += QccIMGImageComponentMSE(&temp1C, &temp2C);
	//	printf(" %f", total);
	total = total/3;
	/*	temp1C.image = SubbandPyramid[0].matrix;
	temp2C.image = SubbandPyramid[2].matrix;
	printf(" 0->2 diff: %f\n",  QccIMGImageComponentMSE(&temp1C, &temp2C));
	temp2C.image = SubbandPyramid[1].matrix;
	printf(" 0->1 diff: %f\n",  QccIMGImageComponentMSE(&temp1C, &temp2C));
	*/
      }
      printf(" %f", total);
    }

    if (OriginalImage.filename[0] != '\0') { // compute ImageMSE with rounding.
      if (numComponents ==1) {
         RoundAndTruncate(&SubbandPyramid[0]);
      }
      else {
      // convert matrix to RGB, round/truncate, then convert back
      //  to given colorspace. 
         ConvertToRGBRoundAndTruncateConvertBack(SubbandPyramid, colortype);
      }

      double total; 
      temp1C.image = OriginalImage.Y.image;
      temp2C.image = SubbandPyramid[0].matrix;
      total = 	 QccIMGImageComponentMSE(&temp1C, &temp2C);
//            printf(" Y%f", total);
      if (numComponents > 1) {
        double tempd;
	temp1C.image = OriginalImage.U.image;
	temp2C.image = SubbandPyramid[1].matrix;
	total += (tempd= QccIMGImageComponentMSE(&temp1C, &temp2C));
//		printf(" U%f", tempd);
	temp1C.image = OriginalImage.V.image;
	temp2C.image = SubbandPyramid[2].matrix;
	total += (tempd = QccIMGImageComponentMSE(&temp1C, &temp2C));
//		printf(" V%f", tempd);
	total = total/3;
	/*	temp1C.image = SubbandPyramid[0].matrix;
	temp2C.image = SubbandPyramid[2].matrix;
	printf(" 0->2 diff: %f\n",  QccIMGImageComponentMSE(&temp1C, &temp2C));
	temp2C.image = SubbandPyramid[1].matrix;
	printf(" 0->1 diff: %f\n",  QccIMGImageComponentMSE(&temp1C, &temp2C));
	*/
      }
      printf(" %f", total);
    }
    
    printf("\n");

  } while (result ==0);


    // write image to a file. 
  OutputImage.Y.num_rows = ImageNumRows;
  OutputImage.Y.num_cols = ImageNumCols;
  OutputImage.Y.image = SubbandPyramid[0].matrix;
  if (numComponents > 1) {
    OutputImage.U.num_rows = ImageNumRows;
    OutputImage.U.num_cols = ImageNumCols;
    OutputImage.U.image = SubbandPyramid[1].matrix;
    OutputImage.V.num_rows = ImageNumRows;
    OutputImage.V.num_cols = ImageNumCols;
    OutputImage.V.image = SubbandPyramid[2].matrix;
  }

  for (i=0; i < numComponents; i++) {
    SubbandComponentFree(&sc[i]);
  }

  if (QccIMGImageSetMaxMin(&OutputImage))
    {
      QccErrorAddMessage("%s: Error calling QccIMGImageSetMaxMin()",
                         argv[0]);
      QccErrorExit();
    }

  if (QccIMGImageWrite(&OutputImage))
    {
      QccErrorAddMessage("%s: Error calling QccIMGImageWrite()",
                         argv[0]);
      QccErrorExit();
    }

  QccWAVWaveletFree(&Wavelet);
  QccIMGImageFree(&OutputImage);
  if (OriginalImage.filename[0] != '\0') {
    QccIMGImageFree(&OriginalImage);
  }
  QccExit;


  return 0;
}
